package com.miui.systemui.tools

import com.google.gson.Gson
import com.google.gson.JsonParser
import java.net.HttpURLConnection
import java.net.URL
import java.nio.charset.StandardCharsets

// 定义成功和失败时的数据类型
typealias SuccessCallback<String> = (String) -> Unit
typealias FailCallback = (String) -> Unit

object AiModel {
    @JvmStatic
    fun requestAi(
        prompt: String = Prompt.promptSystem,
        userString: String,
        success: SuccessCallback<String>,
        fail: FailCallback
    ) {
        try {
            // 创建请求体
            val requestBody = RequestBody(
                "qwen-plus", arrayOf(
                    Message("system", prompt), Message("user", userString)
                ), true
            )

            // 将请求体转换为 JSON
            val gson = Gson()
            val jsonInputString = gson.toJson(requestBody)

            // 创建 URL 对象
            val url = URL("https://dashscope.aliyuncs.com/compatible-mode/v1/chat/completions")
            val httpURLConnection = url.openConnection() as HttpURLConnection

            // 设置请求方法为 POST
            httpURLConnection.requestMethod = "POST"
            httpURLConnection.setRequestProperty("Content-Type", "application/json; utf-8")
            httpURLConnection.setRequestProperty("Accept", "application/json")

            // 若没有配置环境变量，请用百炼API Key将下行替换为：String apiKey = "sk-xxx";
            val apiKey = "sk-d6ed047d7d404ff6b53b6d8985a765bf"
            val auth = "Bearer $apiKey"
            httpURLConnection.setRequestProperty("Authorization", auth)

            // 启用输入输出流
            httpURLConnection.doOutput = true

            httpURLConnection.outputStream.use { os ->
                val input = jsonInputString.toByteArray(StandardCharsets.UTF_8)
                os.write(input, 0, input.size)
            }
            // 获取响应码
            val responseCode = httpURLConnection.responseCode
            println("Response Code: $responseCode")

            // 使用 BufferedReader 读取响应内容
            httpURLConnection.inputStream?.bufferedReader(StandardCharsets.UTF_8).use { reader ->
                reader?.readLines()?.forEach { line ->
                    val content = parseResponse(line)
                    success.invoke(content)
                }
            }
            println("complete+++++++++++++")
        } catch (e: Exception) {
            e.printStackTrace()
            fail.invoke(e.message.toString())
        }
    }


    private fun parseResponse(rawResponse: String): String {
        // 移除 'data: ' 前缀并修剪空白字符
        val jsonString = rawResponse.removePrefix("data: ").trim()
        if (jsonString.isEmpty()) {
            return ""
        }
        if ("[DONE]".equals(jsonString)) {
            return "[DONE]";
        }


        // 使用 JsonParser 解析 JSON 字符串为 JsonObject
        val jsonObject = JsonParser.parseString(jsonString).asJsonObject

        // 获取 "choices" 数组中的第一个对象
        val choices = jsonObject.getAsJsonArray("choices")
        if (choices.size() > 0) {
            for (choice in choices) {
                val delta = choice.asJsonObject.getAsJsonObject("delta")
                val content = delta.get("content")?.asString.orEmpty()
                return content
            }
        }
        return ""
    }

    internal class Message(var role: String, var content: String)

    internal class RequestBody(var model: String, var messages: Array<Message>, var stream: Boolean)
}
